#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <dlfcn.h>
#include <memory.h>
#include <errno.h>
#include <signal.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/mman.h>
#include <sys/time.h>
#include <assert.h>
#include <malloc.h>
#include "list.h"

#define TIME_INTERVAL 2000
#define handle_error(msg) \
	do { perror(msg); exit(EXIT_FAILURE); } while(0)

// the number of time periods elasped
static unsigned long time_period_count = 0;
// the page size
unsigned  page_size;
// the pointer to the log file
static int logfile;
// count of the number of pages per time period
static unsigned long page_count = 0;
// total number of bytes allocated
static unsigned long total_allocated = 0;
// total number of bytes freed
static unsigned long total_freed = 0;

// Pointer to the original free() call
static void (*_free)(void *) = NULL;
// Pointer to the original malloc() call
static void *(*_malloc)(size_t size) = NULL;
// the address returned by the original malloc() call
static void *malloc_res = NULL;

// the data structure that contains the address pointer
// used as node linked to the hash table lists
typedef struct addr_t
{
	void * addr;			//the start pointer of the address
	unsigned long  bytes;		//the bytes allocated
	struct hlist_node list;		//the link list used to link to other nodes
} addr_t;

#define HASH_TABLE_SIZE_BITS 12
#define HASH_TABLE_SIZE (1<<HASH_TABLE_SIZE_BITS)

// the address hash table used to support fast malloc and fast free
struct hlist_head addr_table[HASH_TABLE_SIZE];

// the free wrapper function call
void free(void *addr)
{
//	printf("shared object free get called!\n");

	if(addr == NULL)
		return;

	//initialization
	if(_free == NULL)
	{
		_free = (void(*)(void *))dlsym(RTLD_NEXT, "free");
	}
	
	addr_t * q = NULL;
	struct hlist_node * p = NULL;
	//find the index of the addr in the hash table
	unsigned long index = hash_ptr(addr, HASH_TABLE_SIZE_BITS);
//	printf("free index = %lu\n", index);

	//find the correct addr structure in the hash linked list
	hlist_for_each(p, &addr_table[index])
	{
		q = hlist_entry(p, addr_t, list);
		if(q->addr == addr)
			break;
	}

	if(p)
	{
		total_freed += q->bytes;
		hlist_del_init(p);
		_free(q);
	}	
	else
	{
		printf("*** MY PROGRAM DETECTED ***: double free or free of invalid address: %p ***\n", addr);
	}
	
	_free(addr);
	_free = NULL;
}

//the malloc() wrapper function call
void* malloc(size_t size)
{
//	printf("shared object malloc get called!\n");
	//initialization
	if(_malloc == NULL)
		_malloc = (void*(*)(size_t))dlsym(RTLD_NEXT, "malloc");
	
	
	malloc_res = (void *)memalign(page_size, size);
	if(malloc_res == NULL)
		handle_error("memalign");

	//allocate and init an addr_t structure
	addr_t * aa = (addr_t *)_malloc(sizeof(addr_t));
	if(aa == NULL)
		handle_error("_malloc");
	INIT_HLIST_NODE(&aa->list);
	aa->addr = malloc_res;
	aa->bytes = size;

	//find the hash index in the hash table
	unsigned long index = hash_ptr(aa->addr, HASH_TABLE_SIZE_BITS);
//	printf("malloc index = %lu\n", index);
	hlist_add_head(&aa->list, &addr_table[index]);
		
		
	total_allocated += size;
//	printf("memory allocation address %p\n", malloc_res);
	return malloc_res;
}

static void on_timer(int sig); 

//this function set the timer 
//the SIGVTALRM signal will be sent to the program after micro_secs microseconds 
//which will lead us to the on_timer function
static void set_timer(int micro_secs)
{
        struct itimerval remaining_time;
        remaining_time.it_interval.tv_sec = 0;
        remaining_time.it_interval.tv_usec = micro_secs;
        remaining_time.it_value.tv_sec = 0;
        remaining_time.it_value.tv_usec = micro_secs;

        signal(SIGVTALRM, on_timer);
        setitimer(ITIMER_VIRTUAL, &remaining_time, NULL);
}

//When the time is up, this function will be invoked
static void on_timer(int sig)
{
	char buffer[256];
	sprintf(buffer, "%lu\t\t\t\t%lu\n", time_period_count, page_count);
	if(write(logfile, buffer, strlen(buffer)) != strlen(buffer))
		handle_error("write");
//	printf("time is up!\n");
	page_count = 0;

	addr_t * q = NULL;
	struct hlist_node * p = NULL;


	int i = 0;
	for(; i < HASH_TABLE_SIZE; ++i)
	{
		hlist_for_each(p, &addr_table[i])
		{
			q = hlist_entry(p, addr_t, list);
		
//			printf("addr: %p, bytes: %lu\n", q->addr, q->bytes);
//			void * data1 = (void *)((((unsigned long)q->addr)&(~(page_size - 1))) + page_size);
//			void * data2 = (void *)((((unsigned long)(q->addr - q->bytes))&~(page_size - 1)));
//			unsigned long len = data1 - data2;
//			printf("data1: %p, data2: %p, len: %lu\n", data1, data2, len);

			if(mprotect(q->addr, q->bytes, PROT_NONE) == -1)
				handle_error("mprotect");

		}
	}

	time_period_count++;
        set_timer(TIME_INTERVAL);
        return;
}

static void sig_handler(int sig, siginfo_t * si, void * unused)
{
	printf("Got SIGSEGV at address: 0x%lx\n", (long) si->si_addr);

	void * aligned_addr = (void *)((((unsigned long)si->si_addr)&(~(page_size - 1))));
//	printf("aligned_addr = 0x%lx\n", (long) aligned_addr);

	if(mprotect(aligned_addr, page_size, PROT_READ|PROT_WRITE) == -1)
		handle_error("mprotect");

	++page_count;
}

// the constructor: invoked at the beginning of the program
__attribute__((constructor))
	static void initialize_so() {
		printf("Hi~~~\n");
		
		//init the hash table
		int i = 0;
		for(; i < HASH_TABLE_SIZE; ++i)
			INIT_HLIST_HEAD(&addr_table[i]);

		//bind the SIGSEGV signal
		struct sigaction sa;
		sa.sa_flags = SA_SIGINFO;
		sigemptyset(&sa.sa_mask);
		sa.sa_sigaction = sig_handler;
		if ( sigaction(SIGSEGV, &sa, NULL) == -1)
			handle_error("sigaction");

		//get the page size
		page_size = sysconf(_SC_PAGE_SIZE);
		if(page_size == -1)
			handle_error("sysconf");

		printf("pagesize: %d\n", page_size);
		
		//open log file
		logfile = open("proj3.log", O_RDWR);
		if(logfile < -1)
			handle_error("open");

		const char * str = "#time_period_count\t\t\t\t#page_count\n";
		if(write(logfile, str, strlen(str)) != strlen(str))
			handle_error("write");
		
		//set and start the timer
		set_timer(TIME_INTERVAL);
	}

// the destructor: invoked at the end of the program
__attribute__((destructor))
	static void destroy_so() {
		printf("Goodbye~~~\n");
		printf("total amounts of bytes allocated:\t%lu\n", total_allocated);
		printf("total amounts of bytes freed:\t\t%lu\n", total_freed);
	//	fclose(logfile);
	}
